#!/usr/local/python3/bin/python3.6
import pymongo
from pymongo import UpdateOne
import datetime
from string import digits


def source_mongodb():
    client = pymongo.MongoClient('192.168.1.156', 27017)
    db_auth = client.ipsearch
    db_auth.authenticate("root", "lotut@mongodb")
    db = db_auth.patent
    return db


def target_mongodb():
    client = pymongo.MongoClient('192.168.1.156', 27017)
    db_auth = client.ipsearch
    db_auth.authenticate("root", "lotut@mongodb")
    db = db_auth.patent
    return db


def is_none(data):
    if data is None:
        return True
    return False


# 格式化代理机构
def convert_proxy_org(proxy_org):
    result = {}
    if is_none(proxy_org) is False:
        proxy_org = proxy_org.strip()
        proxy_org_array = proxy_org.split(' ')
        if len(proxy_org_array) == 2:
            if proxy_org_array[0].isdigit() is True:
                result['proxyOrgCode'] = proxy_org_array[0]
                result['proxyOrg'] = proxy_org_array[1]
            else:
                result['proxyOrgCode'] = proxy_org_array[1]
                result['proxyOrg'] = proxy_org_array[0]
        else:
            result['proxyOrg'] = proxy_org_array[0]
    return result


# 格式化地址
def convert_address(address):
    if is_none(address) is False:
        county = address.get('county', None)
        if is_none(county) is False:
            remove_digits_county = str.maketrans('', '', digits)
            res = county.translate(remove_digits_county)
            address['county'] = res
    return address


# 批量更新
def bulk_update(patents):
    target_patent = target_mongodb()
    bulk_update_statement = []
    for patent in patents:
        set_statement = {}
        proxyOrg = patent.get('proxyOrg', None)
        if is_none(proxyOrg) is False:
            set_statement['proxyOrg'] = proxyOrg
        proxyOrgCode = patent.get('proxyOrgCode', None)
        if is_none(proxyOrgCode) is False:
            set_statement['proxyOrgCode'] = proxyOrgCode
        address = patent.get('address', None)
        if is_none(address) is False:
            set_statement['address'] = address
        bulk_update_statement.append(UpdateOne({'_id': {'$eq': patent['_id']}}, {'$set': set_statement}))
    target_patent.bulk_write(bulk_update_statement)


# 解析专利源数据
def parse_patent_source(source):
    proxy_org = convert_proxy_org(source.get('proxyOrg', None))
    source['proxyOrg'] = proxy_org.get('proxyOrg', None)
    source['proxyOrgCode'] = proxy_org.get('proxyOrgCode', None)
    address = convert_address(source.get('address', None))
    source['address'] = address
    return source


# 批量处理查询结果
def batch_process_patent(patent_result):
    patent_array = []
    for re in patent_result:
        patent_array.append(parse_patent_source(re))
        if len(patent_array) == 1000:
            bulk_update(patent_array)
            patent_array = []
    if len(patent_array) > 0:
        bulk_update(patent_array)


# 构造查询语句
def build_query_statement(last_id):
    query_statement = {
        "publishDate": {
            "$gte": datetime.datetime(1985, 1, 1, 16, 0, 0),
            "$lte": datetime.datetime(1999, 12, 31, 16, 0, 0)
        }
    }
    if is_none(last_id) is False:
        query_statement['_id'] = {"$gt": last_id}
    return query_statement


# 仿游标查询
def cursor_find(last_id, limit):
    return list(source_mongodb().find(build_query_statement(last_id)).sort([("_id", pymongo.ASCENDING)]).limit(limit))


if __name__ == "__main__":
    count = 0
    source_patent = source_mongodb()
    print('total: ', source_patent.count(build_query_statement(None)))
    print('begin update')
    result = cursor_find(None, 1000)
    size = len(result)
    if size > 0:
        batch_process_patent(result)
        next_id = result[size - 1].get('_id')
        count = count + size
        print("updated ", count)
        while is_none(next_id) is False:
            result = cursor_find(next_id, 1000)
            size = len(result)
            if size > 0:
                batch_process_patent(result)
                count = count + size
                print("updated ", count)
                next_id = result[size - 1].get('_id')
            else:
                next_id = None
    print('end update')
    print('total updated num ', count)
